#include "DynamicQuantizeLine.h"
#include "alignedface.h"
#include "cartoonnet.h"
#include "facecompose.h"
#include "facedetect.h"
#include <string>
#include <unistd.h>

class Cartoon {
    public:
        Cartoon () {
            face = NULL;
            bgNet = NULL;
            headNet = NULL;
            newWidth = newHeight = 0;
        }
        ~Cartoon () {
            if (face)
                delete face;
            if (bgNet)
                delete bgNet;
            if (headNet)
                delete headNet;
        }
        bool Run (const std::string &infn, const std::string &outfn, const std::string &mode) {
            printf ("file: %s\n", infn.c_str ());
            Mat inImage = imread (infn);
            int srcWidth = inImage.cols;
            int srcHeight = inImage.rows;
            if (inImage.empty ()) {
                std::cerr << "无法加载图像" << std::endl;
                return -1;
            }

            Mat nImage = resizeWithAspectRatio (inImage, 320, 320);
            imwrite ("nImage.jpg", nImage);
            printf ("-------------\n");
            std::string head_model = "model/" + mode + "/cartoon_h.onnx";
            std::string bg_model = "model/" + mode + "/cartoon_bg.onnx";
            Mat         bgMat;

            bgNet = new CartoonNet (bg_model, "bg cartoon");
            bgNet->forward (nImage, bgMat);
            //exit (0);
            // bgMat = nImage;
            if (bgMat.empty ()) {
                return false;
            }
            delete bgNet;
            headNet = new CartoonNet (head_model, "head cartoon");
            bgNet = NULL;
            imwrite ("bg.jpg", bgMat);
            face = new FaceDectect ("model/face/yunet_n_320_320.onnx");
            Mat faceResult;

            face->forward (nImage, faceResult);
            delete face;
            face = NULL;

            if (!faceResult.empty ()) {
                std::cout << faceResult << std::endl;
                for (int i = 0; i < faceResult.rows; i++) {
                    AlignedFace alignedFace;
                    Args        in;

                    if (faceResult.ptr<float> (i)[2] < 48.0f || faceResult.ptr<float> (i)[3] < 48.0f)
                        continue;

                    in.image = nImage;
                    std::vector<float> &vec = in.args;
                    vec.insert (vec.end (), faceResult.ptr<float> (i), faceResult.ptr<float> (i) + faceResult.cols);
                    Args out;
                    alignedFace.forward (in, out);
                    imwrite ("alignedout.jpg", out.image);
                    if (!out.image.empty ()) {
                        Mat inMat, h1Mat;

                        // resize (out.image, inMat, Size (256, 256), 0, 0, INTER_LINEAR);
                        inMat = out.image;
                        {
                            char sname[100];
                            snprintf (sname, 100, "head.in%d.jpg", i);
                            imwrite (sname, out.image);
                        }
                        headNet->forward (inMat, h1Mat);
                        std::cout << "h1Mat: (" << h1Mat.rows << "," << h1Mat.cols << ")"<< std::endl;
                        Mat hMat = h1Mat;
                        // resize (h1Mat, hMat, Size (192, 192), 0, 0, INTER_LINEAR);

                        {
                            char sname[100];
                            snprintf (sname, 100, "head.out%d.jpg", i);
                            imwrite (sname, hMat);
                        }
                        ComposeFace     compose;
                        ComposeFaceArgs cArgs;
                        ComposeFaceArgs cOutArgs;
                        cArgs.transport.create (2, 3, CV_32FC1);
                        for (int y = 0; y < cArgs.transport.rows; y++) {
                            float *f = cArgs.transport.ptr<float> (y);
                            for (int x = 0; x < cArgs.transport.cols; x++) {
                                f[x] = out.args[y * 3 + x];
                            }
                        }
                        // printf("xxxxxxxxxxxx\n");
                        cArgs.dstImage = bgMat;
                        cArgs.srcImage = hMat;
                        compose.forward (cArgs, cOutArgs);
                    }
                }
            }
            Rect roi (0, 0, newWidth, newHeight);
            Mat  cropped = bgMat (roi);
            resize (cropped, bgMat, cv::Size (srcWidth, srcHeight));
            imwrite ("res.jpg", bgMat);
            return true;
        }

    private:
        FaceDectect *face;
        CartoonNet  *bgNet;
        CartoonNet  *headNet;
        Mat          resizeWithAspectRatio (const Mat &image, int width, int height);
        int          newHeight;
        int          newWidth;
};

Mat Cartoon::resizeWithAspectRatio (const Mat &image, int width, int height) {
    if (image.empty () || width <= 0 || height <= 0) {
        return Mat (); // 返回空矩阵
    }

    const float widthRatio = static_cast<float> (width) / image.cols;
    const float heightRatio = static_cast<float> (height) / image.rows;
    const float scaleRatio = std::min (widthRatio, heightRatio);

    newWidth = static_cast<int> (image.cols * scaleRatio);
    newHeight = static_cast<int> (image.rows * scaleRatio);
    printf ("newWidth: %d newHeight: %d\n", newWidth, newHeight);
    const int interpolation = (scaleRatio < 0.5f) ? INTER_AREA : INTER_LINEAR;

    if (image.cols == width && image.rows == height) {
        return image.clone ();
    }

    Mat result (height, width, CV_8UC3, Scalar::all (0));

    if (scaleRatio != 1.0f) {
        Mat resized;
        resize (image, resized, Size (newWidth, newHeight), 0, 0, interpolation);
        resized.copyTo (result (Rect (0, 0, newWidth, newHeight)));
    } else {
        image (Rect (0, 0, std::min (image.cols, width), std::min (image.rows, height)))
            .copyTo (result (Rect (0, 0, std::min (image.cols, width), std::min (image.rows, height))));
    }
    printf ("resule: %d %d\n", result.rows, result.cols);

    return result;
}

int main (int argc, char *argv[]) {
    int         opt;
    std::string infile = "test_face.jpg";
    std::string outfile = "result.jpg";
    std::string mode = "anime";
    int         help_flag = 0;
    while ((opt = getopt (argc, argv, "i:o:m:h")) != -1) {
        switch (opt) {
        case 'i':
            infile = optarg;
            break;
        case 'o':
            outfile = optarg;
            break;
        case 'm':
            mode = optarg;
            break;
        default:
            help_flag = 1;
            break;
        }
    }
    // 显示帮助信息（如果用户指定了 -h 选项）
    if (help_flag) {
        printf ("用法：%s [选项]\n", argv[0]);
        printf ("选项:\n");
        printf ("  -i <文件>    指定输入文件\n");
        printf ("  -o <文件>    指定输出文件\n");
        printf ("  -m <模式>    anime ... 与模型路径名称一致\n");
        printf ("  -h           显示此帮助信息\n");
        return -1;
    }
    CV_DNN_REGISTER_LAYER_CLASS (DynamicQuantizeLinear, DynamicQuantizeLinearLayer);
    Cartoon cartoon;
    cartoon.Run (infile, outfile, mode);
    return 0;
}
// export LD_LIBRARY_PATH=$PWD/../x86/lib
